feat: route arrays_zip on map element types through the codegen dispatcher - #5688

Correctness

Reviewed head 74e92e84e1d99957c514500d5d9cc6db0f4fc311 against base 2da32915c87909d7dbf8d62adb02f8349c36bca0. I found no verified correctness defect in the changed behavior. There is a performance-validation request below.

The maintained Spark 3.5 and 4.0 implementations return null when an input array is null, pad shorter arrays with null fields, and retain the input-derived struct field names. Adding CodegenDispatchFallback lets the existing dispatcher evaluate map-containing ArraysZip trees using Spark's code generation. I checked the recursive map readers and writers as well as the expression semantics. The result keeps non-null struct elements, nullable zipped fields, and the nested map key/value types. The exported Arrow list and map field names agree with Comet's native type serialization.

The existing scalar, array and struct native type gate and convert implementation are unchanged. The per-expression disable switch still acts before dispatch. Disabling spark.comet.exec.scalaUDF.codegen.enabled, rejection by canHandle, or failure to serialize a dispatched input attribute or return type still declines the conversion. A native-path child conversion returning None is not retried through dispatch. The new marker does not alter that path for NullType, and the dispatcher itself rejects NullType outputs.

Local component validation compiled the exact Comet recursive kernel, input and output implementations against Spark 4.0.4, Scala 2.13.16 and Arrow 18.3.0 on JDK 17. All 224 row-and-batch comparisons with Spark expression evaluation passed across seven map shapes with ANSI mode on and off. These covered null and empty arrays, null map elements and values, unequal lengths, nested arrays/structs/maps, and array/struct map keys. The checks included Arrow C Data input/output round trips, exported schema names, and zero outstanding allocator bytes. The probe substitutes the scalar vector adapter and allocator. It does not run the full Comet planner, Rust consumer, JNI query path or SQL fixture suite. Maintained source comparison was limited to Spark 3.5 and 4.0, and no 3.4/4.1 runtime coverage is claimed.

The new SQL queries use table columns and the harness requires Comet execution plus Spark answer equality. They are useful regression coverage, but I have not run them here. At the September 4, 19:52 UTC snapshot, all six head workflow records were action_required, with no head or merge check runs. That is not passing CI evidence.

Performance

Supported native inputs retain the same execution path, so this does not add a per-row dispatcher cost to those cases. Map-containing trees now pay the existing bridge cost: input export/import per batch, expression/kernel cache lookup, Spark row-wise evaluation, and allocation and population of the nested output vector. The cache avoids recompiling the expression for each batch, but zipping maps still materializes result structures and copies their contents into Arrow output buffers.

Keeping surrounding expressions in Comet can avoid a projection fallback, but that benefit needs to be weighed against the bridge and nested-output costs for this newly enabled case. Neither the PR nor the correctness probe measures that tradeoff. The inline P2 asks for a matched microbenchmark before enabling this path by default. This is a performance-validation gap, not a claim that a regression has been measured.

Design

The change uses the existing dispatch decision point and preserves the native implementation for types it already accepts. That keeps map support aligned with Spark without adding another map-specific native implementation or changing the native compatibility gate. The two SQL cases exercise a map result and a mixed projection with scalar arithmetic. The explicit hybrid entry in the expression list reflects this division of execution paths.

Abstraction & complexity

The existing marker is an appropriate extension point for this change. No new dispatcher, map adapter, protocol representation or runtime cache is introduced. Type admission, binding, transport and output allocation remain centralized in the existing implementation. I found no actionable added abstraction or complexity concern in this diff.

CometArraysZip.getSupportLevel still treats NullType as a supported element type through the unchanged isTypeSupported helper, so arrays_zip(array()) by itself stays fully native and never touches the dispatcher. The NullType fallback described in this PR only happens when a NullType array sits next to a MapType array in the same arrays_zip call. In that case getSupportLevel returns Unsupported because of the map argument, CometBatchKernelCodegen.isSupportedDataType then declines the whole bound tree because it doesn't admit NullType, and the entire call, map argument included, falls back to Spark. That matches the pre-PR behavior for that combination, but nothing in arrays_zip.sql exercises it. Could you add a case such as arrays_zip(array(map(k1, v1)), array()) over table columns to confirm the whole expression cleanly falls back instead of partially converting?

The note added to the arrays_zip row in expressions.md reads as if any NullType element causes a fallback on its own, but that's only true in the mixed case above. Would it read better as something like "arrays_zip over a NullType array combined with a map array falls back to Spark", so it doesn't suggest arrays_zip(array()) is affected by itself?

The two new map queries confirm the projection stays in Comet, but they don't cover a null map, a null row, an empty map array, or map arrays of different lengths, the way the existing scalar and struct cases above them already do for their own types. Might be worth adding one of those for the map path too.

Correctness

Rechecked unchanged head 74e92e84 against the new discussion. I found no additional runtime correctness defect. Maintained Spark 3.5/4.0 sources still require a null result for a null input array, longest-array length with null padding, and nullable struct fields inside non-null struct elements. The map route invokes Spark's complete ArraysZip expression. Its descendants retain Spark's own evaluation behavior.

The NullType distinction in the follow-up is supported. CometArraysZip.isTypeSupported accepts NullType recursively, so arrays_zip(array()) remains eligible for the existing native conversion. CometCreateArray serializes the empty array as a literal, and the native zip implementation admits its Null element type. The dispatcher is not involved in that admission decision.

By contrast, arrays_zip(array(map(k1,v1)), array()) first fails the native map gate. The marker enrolls that whole expression in dispatch. After binding, canHandle rejects its NullType-bearing output before payload, input-argument or return-type serialization. The outer call therefore falls back without partially lowering its map child. More generally, dispatch checks recursive result/input types, forbidden expression nodes and the nested-field limit. All input arguments and the return type must then serialize. A decline from the existing Compatible/native converter is still not retried through dispatch.

Could you extend the existing column-backed fixture before merging with a compact map-array table covering null map elements, null map values, null input arrays, empty arrays and unequal map-versus-map lengths. This also exercises reading map-array columns: the two new queries currently construct maps from scalar columns. Retain answer/operator assertions for admitted cases, and add a mixed map/NullType case that asserts clean whole-call fallback. One correction to the coverage discussion: the mixed map/integer query already tests length 1 versus 2 and null padding of the map field. The older null-struct example is commented out, so it is not active coverage.

The current checks show 65 successes, 51 skips and one cancellation. The workflows show four successes, three skips and two cancellations, so this is not an entirely green run set. The successful Spark 3.5 and Spark 4.0 expression logs pass both dictionary variants of arrays_zip.sql. Both checked out merge 1658b1b1 of the assigned base/head and consumed the native artifact ID/digest recorded by that merge's build. The relevant zip/dispatch path and fixture are source-equivalent for these cases. The complete merged tree differs, including a stricter null-short-circuit gate that does not apply to the non-null-intolerant ArraysZip root. The harness disables constant folding and checks answers plus Comet operators. These passes cover the current fixtures, not the missing cases above. The prior 224-comparison probe remains component evidence with substituted scalar adapter/allocation. It was not rerun. Maintained 3.4/4.1 source branches remain unavailable, so this review adds no source-compatibility claim for those versions.

Performance

The existing [P2] benchmark request remains unresolved, with no author reply or measurements. The successful correctness jobs do not quantify the tradeoff between retaining the Comet projection and the dispatcher’s per-row Spark evaluation, nested result allocation and Arrow transfer. Please supply the requested matched map-only and mixed-projection comparison with dispatch enabled/disabled, nonconstant columns, representative batch/map sizes and checked results. This remains an evidence gap, not an observed slowdown. I am not duplicating the inline.

Design

Keep the marker-based route. Correct both the expressions.md note and getUnsupportedReasons: native NullType admission is unchanged, while map-containing calls dispatch only when enabled and accepted by the dispatcher. The current “only ... disabled” reason also omits dispatcher refusal. Update the serde reason source so the normal documentation generation publishes the corrected compatibility pages. The mixed map/NullType case above should pin this distinction. Widening dispatcher type support is unnecessary for this PR.

Abstraction & complexity

The follow-up exposes no need for another abstraction. The existing marker, whole-expression dispatch gate and unchanged native converter already define the boundary. Focus the remaining changes on accurate documentation, durable boundary coverage and the existing performance request. No new independent inline finding is added.

There has been no push since my last round, so the three points from then still stand, and I have two more after another look.

The two new map queries use a plain query, which checks the answer and that the operator stayed in Comet but cannot tell a native expression apart from a codegen-dispatched one. That distinction is the whole point of this change. expect_dispatch(...) and expect_native(...) landed on main in #5610 after this branch was cut, and docs/source/contributor-guide/sql-file-tests.md names CodegenDispatchFallback serdes as exactly the case to annotate. Would you mark the two map queries query expect_dispatch(arrays_zip)?

There is a second half to that. Before the mixin, every plain query in arrays_zip.sql implicitly pinned the native path, because any narrowing of isTypeSupported would have surfaced as a Spark fallback and failed the operator check. With CodegenDispatchFallback in place the same narrowing quietly reroutes those queries through the dispatcher and all of them still pass. So this change costs the file about forty implicit native assertions. Could you put query expect_native(arrays_zip) on at least one column case such as select arrays_zip(a, b) FROM test_arrays_zip? round.sql has the pattern for both halves. One thing to watch on the mixed query: a query carries a single mode, so the map part and the v1 + v2 part need to be separate queries for this to say anything.

The other thing is that the mixin routes more than MapType. getSupportLevel returns Unsupported for every element type isTypeSupported rejects, and the dispatcher takes all of them. Collated strings are in that set on Spark 4.0 and later, because the StringType alternative in isTypeSupported is the default-collation singleton while the dispatcher's isSupportedDataType matches _: StringType, so arrays_zip over an array<string collate UTF8_LCASE> now runs Spark's doGenCode by default where it used to fall back. TimeType is in the same set on Spark 4.1, which is the case the commented-out query at line 204 describes. Both clear supportedDataType and serializeDataType, so both are reachable today with no fixture. array_intersect went through the same widening and got array_intersect_collation.sql behind a MinSparkVersion: 4.0 header. Would the equivalent for arrays_zip make sense, and could the TIME query be uncommented behind MinSparkVersion: 4.1 now that it should work? Since getUnsupportedReasons() is being rewritten anyway, worth phrasing it so it does not read as though MapType is the only case that dispatches.

Last thing is sequencing. #5867 adds the same with CodegenDispatchFallback to CometArraysZip, rewrites the same getUnsupportedReasons(), and edits the same expressions.md row, and #5875 touches arrays.scala and that table too. #5867 also carries a nondeterministic-child gate for ArraysZip that this PR does not, which matters because convert serializes each child twice through the IsNotNull guard. Could we settle an order, or fold the ArraysZip part of one into the other?